import pandas as pd
from pylab import *
import numpy as np

# 设置中文字体和关闭警告
mpl.rcParams['font.family'] = ['Times New Roman', 'SimSun']
mpl.rcParams['axes.unicode_minus'] = False
mpl.rcParams['font.size'] = 14

# 读取数据并转换格式
df = pd.read_csv('rescue_schedule_by_team.csv', encoding='cp1252')
df.columns = ['时间', '1', '2', '3']

# 转换为长格式并清理数据
df_long = df.melt(id_vars='时间', var_name='施工队', value_name='路段').dropna()
df_long['时间'] = df_long['时间'].astype(int)


def merge_intervals(group):
    times = sorted(group['时间'].unique())
    merged = []
    if not times:
        return merged

    current_start = current_end = times[0]
    for t in times[1:]:
        if t == current_end + 1:
            current_end = t
        else:
            merged.append((current_start, current_end))
            current_start = current_end = t
    merged.append((current_start, current_end))
    return merged


# 合并连续时间段
merged = df_long.groupby(['路段', '施工队']).apply(merge_intervals).reset_index()
merged = merged.explode(0).dropna()
merged[['start', 'end']] = pd.DataFrame(merged[0].tolist(), index=merged.index)

# 创建路段排序（按主路段和方向排序）
sections = sorted(df_long['路段'].unique(),
                  key=lambda x: (int(x.split('-')[0]), int(x.split('-')[1])))
y_pos = {section: i + 1 for i, section in enumerate(sections)}

# 创建图形
figure(figsize=(12, len(sections) * 0.7))  # 高度根据路段数量调整
xlim(24, 96)
xticks(np.arange(24, 97, 12))
ylim(0.5, len(sections) + 0.5)
yticks(range(1, len(sections) + 1), sections)

# 绘制每个时间段
for _, row in merged.iterrows():
    y = y_pos[row['路段']]
    duration = row['end'] - row['start'] + 1
    barh(y, duration, left=row['start'],
         color='white', edgecolor='black', height=0.3)
    text(row['start'] + duration / 2, y, row['施工队'],
         ha='center', va='center', fontsize=12)

# 图表美化
xlabel('时间/h', fontsize=14)
ylabel('受损路段编号', fontsize=14)
gca().invert_yaxis()  # 让第一个路段显示在顶部
grid(True, linestyle='--', alpha=0.5)
gca().spines['top'].set_visible(False)
gca().spines['right'].set_visible(False)
tight_layout()
savefig('road_repair_schedule.png', dpi=300, bbox_inches='tight')
show()